import open3d as o3d
import numpy as np
import cv2

'''取某个位置上第多少个z点进行投影
将点云筛选并映射成栅格地图'''


# 提取并保存点云中的最小到第十小的数据
def get_2d_point(cloud_path, output_path, visualize=False):
    pcd = o3d.io.read_point_cloud(cloud_path)
    points = np.asarray(pcd.points)
    colors = np.asarray(pcd.colors)

    # 创建一个字典用于存储每个 (X, Y) 坐标对应的 z 值列表及颜色
    point_groups = {}

    # 遍历点云数据
    for i in range(points.shape[0]):
        x, y, z = points[i]

        key = (int(x * 10), int(y * 10))  # 使用整数作为键，乘以 10 是为了提高分辨率
        # 如果字典中不存在当前的 key，初始化空列表
        if key not in point_groups:
            point_groups[key] = []
        # 将当前点的 z 值和颜色附加到对应的 key 的列表中
        point_groups[key].append((z, colors[i]))

    projected_points = []
    projected_colors = []

    # 遍历所有的 (X, Y) 组
    for key, values in point_groups.items():
        # 按照 z 值排序
        values.sort(key=lambda x: x[0])

        # todo  两种选择：保存所有还是某个点
        # # 保留最小 z 值到第 10 小 z 值之间的所有点
        for z_val, color in values[:8]:  # 如果少于10个点，则取全部
            x, y = key[0] / 10.0, key[1] / 10.0  # 还原坐标
            projected_points.append([x, y, z_val])  # 保留原始的 z 值
            projected_colors.append(color)
        # 获取第 10 个点的颜色，如果少于 10 个点，跳过
        # if len(values) >= 10:
        #     tenth_point_color = values[9][1]  # 获取第 10 个点的颜色
        #     # x, y = key[0] / 10.0, key[1] / 10.0  # 还原坐标
        #     projected_points.append([x, y, 0])  # z 值设为 0，因为是在 XOY 平面上
        #     projected_colors.append(tenth_point_color)

    # 转换为 numpy 数组
    projected_points = np.array(projected_points)
    projected_colors = np.array(projected_colors)

    #  点云创建+保存+可视化
    projected_pcd = o3d.geometry.PointCloud()
    projected_pcd.points = o3d.utility.Vector3dVector(projected_points)
    projected_pcd.colors = o3d.utility.Vector3dVector(projected_colors)
    o3d.io.write_point_cloud(output_path, projected_pcd)
    print(f"Filtered point cloud saved to {output_path}")
    if visualize:
        o3d.visualization.draw_geometries([projected_pcd])

    return projected_points, projected_colors


def replace_colors(points, colors, target_color_list, new_color_list):
    # 将颜色值转换为0-1之间的值
    target_colors = np.array(target_color_list) / 255.0
    new_colors = np.array(new_color_list) / 255.0

    for target_color, new_color in zip(target_colors, new_colors):
        # 找到所有颜色匹配的点
        mask = np.all(colors == target_color, axis=1)
        # 替换颜色
        colors[mask] = new_color

    return points, colors


# 绘制2D栅格地图并保存映射关系
# 绘制2D栅格地图并保存映射关系
def plot_2d_map_with_lookup(projected_points, projected_colors, color_image_path, grey_image_path, scale_factor=10):
    projected_points_2d = projected_points[:, :2]
    projected_points_2d *= scale_factor

    # 计算图像边界和尺寸
    min_bounds = np.min(projected_points_2d, axis=0)
    max_bounds = np.max(projected_points_2d, axis=0)
    image_size = (max_bounds - min_bounds + 1).astype(int)  # 图像的宽度和高度

    image = np.ones((image_size[1], image_size[0], 3), dtype=np.uint8) * np.array([128, 128, 128], dtype=np.uint8)

    pixel_to_pointcloud = {}
    max_z_dict = {}

    # 遍历所有的投影点，计算每个像素点对应的颜色和坐标
    for i in range(projected_points_2d.shape[0]):
        x, y = projected_points_2d[i] - min_bounds
        x, y = int(x), int(y)
        y = image_size[1] - y - 1
        z = projected_points[i][2]

        # 保存最高点云的像素点
        if (x, y) not in max_z_dict or z > max_z_dict[(x, y)][0]:
            max_z_dict[(x, y)] = (z, projected_colors[i])
            pixel_to_pointcloud[(x, y)] = {'coordinate': projected_points[i], 'color': projected_colors[i]}

    # 根据最高点云颜色绘制图像
    for (x, y), (z, color) in max_z_dict.items():
        image[y, x] = (color * 255).astype(np.uint8)

    # 在图像上标记 (0, 0) 点
    # origin_x, origin_y = -min_bounds  # 假设原点在 (0, 0)
    # origin_x, origin_y = int(origin_x), int(origin_y)
    # if 0 <= origin_x < image_size[0] and 0 <= origin_y < image_size[1]:
    #     cv2.circle(image, (origin_y, origin_x), radius=5, color=(0, 0, 255), thickness=-1)  # 红色点

    # 保存彩色图像和灰度图像
    cv2.imwrite(color_image_path, image)
    print(f"Color image saved as {color_image_path}")
    grey_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    cv2.imwrite(grey_image_path, grey_image)
    print(f"Grey image saved as {grey_image_path}")

    apply_filters(color_image_path, grey_image_path)

    # 调用保存像素-点云映射的函数，并传入计算出的图像大小
    complete_pixel_to_pointcloud = save_all_pixel_to_pointcloud_mapping(pixel_to_pointcloud,
                                                                        "data/pixel_to_pointcloud.txt", image_size)

    return pixel_to_pointcloud, image_size, complete_pixel_to_pointcloud


# 保存原始地图中的所有像素位置、点云坐标和颜色信息（包括无点云的像素点）
def save_all_pixel_to_pointcloud_mapping(mapping, file_path, image_size):
    # 创建一个新字典，用来存储有点云和无点云的像素
    complete_pixel_to_pointcloud = {}

    # 遍历所有像素位置
    for x in range(image_size[0]):
        for y in range(image_size[1]):
            if (x, y) in mapping:
                # 如果有点云数据，存储坐标和颜色
                complete_pixel_to_pointcloud[(x, y)] = mapping[(x, y)]
            else:
                # 如果没有点云数据，坐标和颜色都为 None
                complete_pixel_to_pointcloud[(x, y)] = {'coordinate': (None, None, None), 'color': (None, None, None)}

    # 将所有像素位置、点云坐标和颜色信息保存到文件中
    with open(file_path, 'w') as f:
        for x in range(image_size[0]):
            for y in range(image_size[1]):
                coord = complete_pixel_to_pointcloud[(x, y)]['coordinate']
                color = complete_pixel_to_pointcloud[(x, y)]['color']
                f.write(f"{x},{y}:{coord[0]},{coord[1]},{coord[2]}:{color[0]},{color[1]},{color[2]}\n")

    print(f"Pixel to point cloud mapping (with all pixels) saved to {file_path}")
    return complete_pixel_to_pointcloud


# 保存经过滤波处理后的像素位置、点云坐标和颜色信息
def save_filtered_pixel_to_pointcloud_mapping(image_path, pixel_to_pointcloud, suffix, img_size):
    filtered_pixel_to_pointcloud = {}
    # 读取经过滤波处理后的图像
    filtered_image = cv2.imread(image_path)
    # 将滤波处理后的图像映射到点云数据    将滤波前的数值替换为滤波后的数值
    for pixel, data in pixel_to_pointcloud.items():
        x, y = pixel
        if 0 <= x < filtered_image.shape[1] and 0 <= y < filtered_image.shape[0]:
            color = filtered_image[y, x]
            color = color / 255.0  # 将颜色值转换为0-1之间的值
            filtered_pixel_to_pointcloud[pixel] = {
                'coordinate': data['coordinate'],
                'color': color
            }
    # 保存到文件
    mapping_file_path = f'data/pixel_to_pointcloud_mapping_{suffix}.txt'
    save_all_pixel_to_pointcloud_mapping(filtered_pixel_to_pointcloud, mapping_file_path, img_size)


# 对生成的2d栅格图像进行中值滤波和双边滤波处理+两种方法融合处理
def apply_filters(color_image_path, grey_image_path):
    color_image = cv2.imread(color_image_path)
    grey_image = cv2.imread(grey_image_path, cv2.IMREAD_GRAYSCALE)

    # 中值滤波处理
    color_image_median = cv2.medianBlur(color_image, ksize=5)
    grey_image_median = cv2.medianBlur(grey_image, ksize=5)

    # 双边滤波处理
    color_image_bilateral = cv2.bilateralFilter(color_image, d=9, sigmaColor=50, sigmaSpace=50)
    grey_image_bilateral = cv2.bilateralFilter(grey_image, d=9, sigmaColor=50, sigmaSpace=50)

    # 先中值后双边
    color_image_median_bilateral = cv2.bilateralFilter(color_image_median, d=9, sigmaColor=50, sigmaSpace=50)
    grey_image_median_bilateral = cv2.bilateralFilter(grey_image_median, d=9, sigmaColor=50, sigmaSpace=50)

    # 保存去噪后的图像
    cv2.imwrite(color_image_path.replace('.png', '_median.png'), color_image_median)
    cv2.imwrite(grey_image_path.replace('.png', '_median.png'), grey_image_median)
    # cv2.imwrite(color_image_path.replace('.png', '_bilateral.png'), color_image_bilateral)
    # cv2.imwrite(grey_image_path.replace('.png', '_bilateral.png'), grey_image_bilateral)
    # cv2.imwrite(color_image_path.replace('.png', '_median_bilateral.png'), color_image_median_bilateral)
    # cv2.imwrite(grey_image_path.replace('.png', '_median_bilateral.png'), grey_image_median_bilateral)

    print(f"Filtered images saved with '_median' and '_bilateral' and '_median_bilateral' suffixes.")


if __name__ == '__main__':
    cloud_path = 'data/save_cloud1.pcd'
    output_cloud_path = 'data/save_cloud1_10.pcd'
    color_image_path = 'data/color_image_10.png'  # 存储生成图片的路径
    grey_image_path = 'data/grey_image_10.pgm'
    mapping_file_path = 'data/pixel_to_pointcloud_mapping.txt'
    # 把道路与人行道置换成白色；未建图地区置换成灰色
    target_color_list = [[128, 64, 128], [0, 0, 192], (64, 0, 128)]
    new_color_list = [[255, 255, 255], [255, 255, 255], [255, 255, 255]]
    # 获取2D点并保存
    projected_points, projected_colors = get_2d_point(cloud_path, output_cloud_path)
    new_points, new_colors = replace_colors(projected_points, projected_colors, target_color_list, new_color_list)

    pixel_to_pointcloud, image_size, complete_pixel_to_pointcloud = plot_2d_map_with_lookup(new_points, new_colors,
                                                                                            color_image_path,
                                                                                            grey_image_path)

    # save_all_pixel_to_pointcloud_mapping(pixel_to_pointcloud, mapping_file_path)
    save_filtered_pixel_to_pointcloud_mapping('data/color_image_10_median.png', complete_pixel_to_pointcloud, 'median',
                                              image_size)
